import io

import numpy as np

# imports
import pytest
from matplotlib.colors import Colormap
from PIL import Image

# function to test (see above for full class definition, including _repr_png_)

# ---- UNIT TESTS FOR Colormap._repr_png_ ----


# Helper: Minimal subclass to provide a _lut and _init for testing
class DummyColormap(Colormap):
    def _init(self):
        # Simple grayscale LUT for testing: N+3 rows (for under, over, bad)
        self._lut = np.zeros((self.N + 3, 4), dtype=float)
        for i in range(self.N):
            gray = i / (self.N - 1) if self.N > 1 else 0.0
            self._lut[i] = (gray, gray, gray, 1.0)
        # Under/over colors
        self._lut[self.N] = (1.0, 0.0, 0.0, 1.0)  # under: red
        self._lut[self.N + 1] = (0.0, 1.0, 0.0, 1.0)  # over: green
        self._lut[self.N + 2] = (0.0, 0.0, 0.0, 0.0)  # bad: transparent
        self._isinit = True


@pytest.fixture
def colormap_256():
    # Standard 256-level grayscale colormap
    return DummyColormap("dummy_gray", N=256)


@pytest.fixture
def colormap_small():
    # Small colormap for edge tests
    return DummyColormap("tiny", N=2)


@pytest.fixture
def colormap_large():
    # Large colormap for large scale tests
    return DummyColormap("large", N=999)


# --- BASIC TEST CASES ---


def test_png_basic_validity(colormap_256):
    """Test that _repr_png_ returns valid PNG bytes that can be opened as an image."""
    codeflash_output = colormap_256._repr_png_()
    png_bytes = codeflash_output  # 1.05ms -> 1.01ms (4.12% faster)
    # Try to open as PNG image
    img = Image.open(io.BytesIO(png_bytes))


def test_png_metadata(colormap_256):
    """Test that PNG metadata includes expected fields and values."""
    codeflash_output = colormap_256._repr_png_()
    png_bytes = codeflash_output  # 1.01ms -> 951μs (6.30% faster)
    img = Image.open(io.BytesIO(png_bytes))
    meta = img.info
    # Author and Software should contain matplotlib version and url
    author = meta.get("Author")


def test_png_different_names():
    """Test that the PNG metadata Title/Description reflect the colormap name."""
    cmap = DummyColormap("foobar", N=10)
    codeflash_output = cmap._repr_png_()
    png_bytes = codeflash_output  # 906μs -> 881μs (2.78% faster)
    img = Image.open(io.BytesIO(png_bytes))
    meta = img.info


def test_png_basic_image_content(colormap_256):
    """Test that the PNG image visually encodes the colormap: left pixel is black, right is white."""
    codeflash_output = colormap_256._repr_png_()
    png_bytes = codeflash_output  # 961μs -> 945μs (1.78% faster)
    img = Image.open(io.BytesIO(png_bytes))
    # Convert to numpy array for pixel checks
    arr = np.array(img)


# --- EDGE TEST CASES ---


def test_png_tiny_colormap(colormap_small):
    """Test with a colormap of only 2 colors (N=2)."""
    codeflash_output = colormap_small._repr_png_()
    png_bytes = codeflash_output  # 888μs -> 821μs (8.09% faster)
    img = Image.open(io.BytesIO(png_bytes))
    arr = np.array(img)
    # Only two unique RGB colors should appear (black and white)
    unique_colors = set(tuple(rgba[:3]) for row in arr for rgba in row)


def test_png_empty_colormap():
    """Test edge case: N=1 (degenerate colormap)."""
    cmap = DummyColormap("empty", N=1)
    codeflash_output = cmap._repr_png_()
    png_bytes = codeflash_output  # 881μs -> 811μs (8.61% faster)
    img = Image.open(io.BytesIO(png_bytes))
    arr = np.array(img)
    # All pixels should be the same color (since only 1 color in LUT)
    first_pixel = tuple(arr[0, 0])


def test_png_non_ascii_name():
    """Test that non-ASCII colormap names are handled in metadata."""
    name = "cmap_μλΩ"
    cmap = DummyColormap(name, N=10)
    codeflash_output = cmap._repr_png_()
    png_bytes = codeflash_output  # 897μs -> 835μs (7.35% faster)
    img = Image.open(io.BytesIO(png_bytes))
    meta = img.info


def test_png_alpha_channel(colormap_256):
    """Test that alpha channel is always 255 for all pixels (fully opaque)."""
    codeflash_output = colormap_256._repr_png_()
    png_bytes = codeflash_output  # 959μs -> 926μs (3.58% faster)
    img = Image.open(io.BytesIO(png_bytes))
    arr = np.array(img)


def test_png_handles_large_N():
    """Test that a large N colormap (close to 1000) works and is performant."""
    cmap = DummyColormap("largeN", N=999)
    codeflash_output = cmap._repr_png_()
    png_bytes = codeflash_output  # 1.26ms -> 1.23ms (2.28% faster)
    img = Image.open(io.BytesIO(png_bytes))


def test_png_handles_zero_N():
    """Test edge case: N=0 should produce an image with all 'under' color."""
    cmap = DummyColormap("zero", N=0)
    codeflash_output = cmap._repr_png_()
    png_bytes = codeflash_output  # 929μs -> 893μs (4.04% faster)
    img = Image.open(io.BytesIO(png_bytes))
    arr = np.array(img)


def test_png_performance_large_colormap(colormap_large):
    """Test that _repr_png_ completes in reasonable time for large LUTs."""
    import time

    start = time.time()
    codeflash_output = colormap_large._repr_png_()
    png_bytes = codeflash_output  # 1.29ms -> 1.20ms (8.18% faster)
    elapsed = time.time() - start
    # Check image is valid
    img = Image.open(io.BytesIO(png_bytes))


def test_png_large_image_content(colormap_large):
    """Test that the leftmost and rightmost columns of a large LUT colormap are correct."""
    codeflash_output = colormap_large._repr_png_()
    png_bytes = codeflash_output  # 1.26ms -> 1.21ms (3.40% faster)
    img = Image.open(io.BytesIO(png_bytes))
    arr = np.array(img)


def test_png_unique_colors_in_row(colormap_large):
    """Test that a row in the PNG image contains many unique colors for large N."""
    codeflash_output = colormap_large._repr_png_()
    png_bytes = codeflash_output  # 1.25ms -> 1.21ms (3.07% faster)
    img = Image.open(io.BytesIO(png_bytes))
    arr = np.array(img)
    # Count unique colors in the first row (should be >500 for 999-level grayscale)
    unique_colors = {tuple(rgb) for rgb in arr[0, :, :3]}


# --- DETERMINISM/REPEATABILITY ---


def test_png_deterministic_output(colormap_256):
    """Test that repeated calls to _repr_png_ produce identical PNG bytes."""
    codeflash_output = colormap_256._repr_png_()
    png1 = codeflash_output  # 963μs -> 918μs (4.88% faster)
    codeflash_output = colormap_256._repr_png_()
    png2 = codeflash_output  # 785μs -> 727μs (7.98% faster)


def test_png_deterministic_image_content(colormap_256):
    """Test that image content is the same between runs."""
    img1 = Image.open(
        io.BytesIO(colormap_256._repr_png_())
    )  # 987μs -> 943μs (4.65% faster)
    img2 = Image.open(
        io.BytesIO(colormap_256._repr_png_())
    )  # 780μs -> 713μs (9.42% faster)
    arr1 = np.array(img1)
    arr2 = np.array(img2)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.